package org.budo.apache.sharding.algorithm;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.shardingsphere.spi.masterslave.MasterSlaveLoadBalanceAlgorithm;
import org.budo.sharding.hint.aop.AnnotationRoutingDataSourceAopInterceptor;
import org.budo.support.lang.util.ListUtil;

import lombok.Getter;
import lombok.Setter;
import lombok.experimental.Accessors;

/**
 * @author lmw
 */
@Getter
@Setter
@Accessors(chain = true)
public class ApacheShardingWeightedRoundRobinMasterSlaveLoadBalanceAlgorithm implements MasterSlaveLoadBalanceAlgorithm {
    private static final Map<String, AtomicInteger> _COUNT_MAP = new ConcurrentHashMap<String, AtomicInteger>();

    private static final Map<List<String>, List<String>> _WEIGHTED_SLAVE_DATA_SOURCE_NAMES_CACHE = new ConcurrentHashMap<List<String>, List<String>>();

    public static Map<String, Integer> _SLAVE_DATA_SOURCE_WEIGHT = new HashMap<String, Integer>();

    private Properties properties;

    @Override
    public String getDataSource(String logicDataSourceName, String masterDataSourceName, List<String> slaveDataSourceNames) {
        String annotatedRoutingDataSource = AnnotationRoutingDataSourceAopInterceptor._ROUTING_DATA_SOURCE_THREAD_LOCAL.get();
        if (null != annotatedRoutingDataSource) { // 直接通过注解指定当前走哪个库
            return annotatedRoutingDataSource;
        }

        // 计数器
        AtomicInteger count = _COUNT_MAP.get(logicDataSourceName);
        if (null == count) {
            count = new AtomicInteger(0);
            _COUNT_MAP.put(logicDataSourceName, count);
        }

        // 加权后的列表
        List<String> weightedSlaveDataSourceNames = this.getWeightedSlaveDataSourceNames(slaveDataSourceNames);

        // 到顶归零
        int max = weightedSlaveDataSourceNames.size();
        count.compareAndSet(max, 0); // 如果 count = max，就将它修改为0

        int index = count.getAndAdd(1) % max; // 返回除 max 余数，然后 +1
        return weightedSlaveDataSourceNames.get(index);
    }

    public static void clearWeightCache() {
        _WEIGHTED_SLAVE_DATA_SOURCE_NAMES_CACHE.clear();
    }

    /**
     * 加权后的列表
     */
    private List<String> getWeightedSlaveDataSourceNames(List<String> slaveDataSourceNames) {
        List<String> weightedNames = _WEIGHTED_SLAVE_DATA_SOURCE_NAMES_CACHE.get(slaveDataSourceNames);
        if (null != weightedNames) {
            return weightedNames;
        }

        weightedNames = this.buildWeightedSlaveDataSourceNames(slaveDataSourceNames);
        _WEIGHTED_SLAVE_DATA_SOURCE_NAMES_CACHE.put(slaveDataSourceNames, weightedNames);
        return weightedNames;
    }

    public List<String> buildWeightedSlaveDataSourceNames(List<String> slaveDataSourceNames) {
        if (null == _SLAVE_DATA_SOURCE_WEIGHT || _SLAVE_DATA_SOURCE_WEIGHT.size() <= 1) { // 如果只有1个，就不处理了
            return slaveDataSourceNames;
        }

        List<String> list = new ArrayList<String>();
        for (String each : slaveDataSourceNames) {
            Integer weight = _SLAVE_DATA_SOURCE_WEIGHT.get(each);
            if (null == weight) { // 未配置权重，默认为1
                list.add(each);
                continue;
            }

            // 权重小于1
            if (weight < 1) {
                continue;
            }

            // 指定权重
            list.addAll(ListUtil.duplicate(each, weight));
        }

        if (null == list || list.size() <= 0) {
            return slaveDataSourceNames; // 如果全部配成0就全部1:1
        }

        // 乱序,尽量更平均
        Collections.shuffle(list);
        return list;
    }

    public List<String> dataSourceNames() {
        if (null == _SLAVE_DATA_SOURCE_WEIGHT || _SLAVE_DATA_SOURCE_WEIGHT.isEmpty()) {
            return Collections.EMPTY_LIST;
        }

        Set<String> keySet = _SLAVE_DATA_SOURCE_WEIGHT.keySet();
        return new ArrayList<String>(keySet);
    }

    @Override
    public String getType() {
        return ApacheShardingWeightedRoundRobinMasterSlaveLoadBalanceAlgorithm.class.getName();
    }

    @Override
    public void setProperties(Properties properties) {
        this.properties=properties;
    }
}
